import numpy as np
import cv2

def getMaskRed(img):
    """
    提取图中红色部分
    """
    #转化为hsv空间
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    print(hsv.shape)

    #颜色在hsv空间下的上下限156-180和0-10
    low_hsv = np.array([156,43,46])
    high_hsv = np.array([180,255,255])

    #使用opencv的inRange函数提取颜色
    maskRed1 = cv2.inRange(hsv,lowerb=low_hsv,upperb=high_hsv)
    maskRed2 = cv2.inRange(hsv,lowerb=np.array([0,43,46]),upperb=np.array([10,255,255]))
    maskRed = cv2.bitwise_or(maskRed1,maskRed2)
    #redImg = cv2.bitwise_and(img,img,mask=maskRed)
    return maskRed;

def getMaskGreen(img):
    """
    提取图中绿色部分
    """
    #转化为hsv空间
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    print(hsv.shape)

    #颜色在hsv空间下的上下限156-180和0-10
    low_hsv = np.array([35,43,46])
    high_hsv = np.array([77,255,255])

    #使用opencv的inRange函数提取颜色
    maskGreen = cv2.inRange(hsv,lowerb=low_hsv,upperb=high_hsv)
    #colorImg = cv2.bitwise_and(img,img,mask=maskGreen)
    return maskGreen;


if __name__ == '__main__':
    src = r"D:\data\CUGW\test4.bak\g2\IMG_5355.JPG"
    img = cv2.imread(src)
    img = cv2.GaussianBlur(img, (7, 7), 0)
    #红色提取出来
    redMask = getMaskRed(img)
    greenMask = getMaskGreen(img)
    maskColor = cv2.bitwise_or(redMask,greenMask);
    colorImg = cv2.bitwise_and(img,img,mask=maskColor)
    cv2.imwrite("maskColor.jpg",colorImg);
    cv2.imshow("maskColor", colorImg)
    cv2.waitKey()